import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.transform import Rotation as scipy_R
from typing import Union, List
from shapely.geometry import LineString, Polygon
import pickle
import sys
sys.path.pop(-2)
from mmdet3d.ops.voxel_interpolation import voxel_interpolation_cpu as VI

# sys.path.pop(0)
# with open('../tmp/result/result_aux_470.pkl','rb') as f:
#     data = pickle.load(f)

def init_axes(dpi: float=100) -> plt.Axes:
    fig = plt.figure(dpi=dpi)
    return fig.add_axes([0, 0, 1, 1])

def matplot_2d_cloud(cloud_x, cloud_y, scale=0.1, rgb='b',
                     alpha=1, label=None, axes: plt.Axes=None) -> None:
    """ plot 2D cloud with matplotlib
    :param cloud_x: list of cloud x
    :param cloud_y: list of cloud y
    :param scale: size of cloud points, float or list of float
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param label: label of cloud
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes(dpi=1000)
    if label is None:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none')
    else:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes
def matplot_polygon(points_x, points_y, scale=0.1, rgb='g',
                    alpha=1, axes:plt.Axes=None) -> None:
    if axes is None:
        axes = init_axes()
    assert len(points_x) == len(points_y) and len(points_x) > 1
    x_arr = np.append(np.array(points_x), points_x[0])
    y_arr = np.append(np.array(points_y), points_y[0])
    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes
def matplot_lines(points_x, points_y, scale=0.1, rgb='g',
                    alpha=1, axes:plt.Axes=None) -> None:
    if axes is None:
        axes = init_axes()
    assert len(points_x) == len(points_y) and len(points_x) > 1
    x_arr = np.array(points_x)
    y_arr = np.array(points_y)
    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes


DEFAULT_ALL_CLOSE_PRECISION = 1.e-4

def euler_to_rotation_matrix(euler_xyz: Union[np.ndarray, list], degree=False) -> np.ndarray:
    return scipy_R.from_euler('XYZ', euler_xyz, degrees=degree).as_matrix()

def bounding_box_to_points(position: list, size: list, rotation: list) -> np.ndarray:
    """ calculate 3D bounding box's points
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    """
    assert len(position) == len(size) == len(rotation) == 3
    pos_arr = np.array(position)
    size_arr = np.array(size)
    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],
                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2
    return euler_to_rotation_matrix(rotation).dot((pts_prefix * size_arr).T).T + pos_arr
def matplot_bounding_box(position: list, size: list, rotation: list, scale=0.1,
                         rgb='k', alpha=1, axes: plt.Axes=None) -> None:
    """ plot 3D bounding box with matplotlib
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    :param scale: line width of box
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes()
    box_pts = bounding_box_to_points(position, size, rotation)[:5, :2]
    axes.plot(box_pts[:, 0], box_pts[:, 1], c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')
    return axes
def draw_multi_grid_range(axes, grid_size=0.8, bv_range=(-51.2, -51.2, 51.2, 51.2),rgb='k'):
    Lx, Ly, Lz, Lx1, Ly1, Lz1 = [], [], [], [], [], []

    Lx = np.arange(bv_range[0], bv_range[2] + 0.1, grid_size) #bugfix
    Ly1 = np.arange(bv_range[1], bv_range[3] + 0.1, grid_size)
    Lx = np.stack([Lx, Lx], axis=0).transpose().reshape(-1)
    Ly1 = np.stack([Ly1, Ly1], axis=0).transpose().reshape(-1)

    for i in range(int(Lx.shape[0] / 2)):
        if (i + 1) % 2 == 0:
            Ly.append(bv_range[1])
            Ly.append(bv_range[1 + 2])
        else:
            Ly.append(bv_range[1 + 2])
            Ly.append(bv_range[1])
    Ly = np.array(Ly)
    Lz = np.zeros(Ly.shape[0])

    for i in range(int(Ly1.shape[0] / 2)):
        if (i + 1) % 2 == 0:
            Lx1.append(bv_range[0])
            Lx1.append(bv_range[0 + 2])
        else:
            Lx1.append(bv_range[0 + 2])
            Lx1.append(bv_range[0])
    # print(Lx)
    Lx1 = np.array(Lx1)
    Lz1 = np.zeros(Lx1.shape[0])

    axes = matplot_lines(Lx1, Ly1, rgb=rgb, scale=0.05, axes=axes)
    axes = matplot_lines(Lx, Ly, rgb=rgb, scale=0.05, axes=axes)
    # Lx = np.concatenate([Lx,L])
    return axes
def in_range(gt_dense_pts, pts_range):
    pts_mask = np.logical_and(gt_dense_pts[:, 1] >= pts_range[0], gt_dense_pts[:, 1] <= pts_range[3])
    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] <= pts_range[4])
    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] >= pts_range[1])

    return pts_mask

def outlier_gt_filter(gt_dense_pts, pts_range):
    lines_num = int(np.max(gt_dense_pts[:, 0]) + 1)

    pts_mask = in_range(gt_dense_pts, pts_range)
    new_gt_dense_pts = np.zeros(gt_dense_pts.shape, dtype=np.float32)

    id_count = 0
    filled_size = 0
    for i in range(lines_num):
        pos_mask = (gt_dense_pts[:, 0] == i)
        if (pts_mask[pos_mask].sum() == 0):
            continue

        lines_gt = gt_dense_pts[pos_mask, :]
        lines_gt[:, 0] = id_count
        id_count += 1

        circle = (np.abs(lines_gt[0, 1] - lines_gt[-1, 1])**2 + np.abs(lines_gt[0, 2] - lines_gt[-1, 2])**2) < 1
        if(circle and
            (pts_mask[pos_mask].sum() != pos_mask.sum()) and
            in_range(lines_gt[0, :].reshape(1, -1), pts_range)[0]):     #环形 且 存在越界部分 且 起点在范围内
            first_outside = np.where(~pts_mask[pos_mask])[0][0]
            lines_gt = np.concatenate([lines_gt[first_outside: -1, :],
                                        lines_gt[:first_outside, :],
                                        lines_gt[first_outside, :].reshape(1, -1)])

        new_gt_dense_pts[filled_size: filled_size + len(lines_gt), :] = lines_gt
        filled_size += len(lines_gt)

    new_gt_dense_pts = new_gt_dense_pts[:filled_size]
    return new_gt_dense_pts
def decode_one_voxels(dt_pts):
    dx = (dt_pts <= 1) * 0 + \
            ((1 < dt_pts) & (dt_pts <= 2)) * (dt_pts - 1) + \
            ((2 < dt_pts) & (dt_pts <= 3)) * 1 + \
            (3 < dt_pts) * (4 - dt_pts)
    dy = (dt_pts <= 1) * dt_pts + \
            ((1 < dt_pts) & (dt_pts <= 2)) * 1 + \
            ((2 < dt_pts) & (dt_pts <= 3)) * (3 - dt_pts) + \
            (3 < dt_pts) * 0
    return dx, dy

def test():
    with open('./tmp/result/result_aux_map_gt.pkl','rb') as f:
        data = pickle.load(f)
    frame = np.random.randint(low=0,high=len(data))
    # import sys
    # sys.path.append("..")
    # from mmdet3d.ops.voxel_interpolation import voxel_interpolation_cpu as VI

    rgb = ['r','y','b']
    gt_maps = data[frame]['gt_maps_3d'][0][0]
    max_objs = 10000
    # pc_range = torch.tensor(self.train_cfg['point_cloud_range'])
    # voxel_size = torch.tensor(self.train_cfg['voxel_size'])

    # feature_map_size = grid_size[:2] // self.train_cfg['out_size_factor']
    vs = [0.8, 0.8, 8]
    pcr = [-51.2, -51.2, -5.0, 51.2, 51.2, 3.0]
    gt_pts = gt_maps.cpu().numpy()
    gt_pts = gt_pts.astype(np.float32)
    gt_pts = outlier_gt_filter(gt_pts, pcr)
    dense_pts = np.zeros((max_objs, 8), dtype=np.float32)
    import time
    st = time.time()
    size = VI.gt_lines_voxel_interpolation(gt_pts, vs, pcr, 
                                        max_objs, dense_pts)
    dense_pts = dense_pts[:size, :]
    print(time.time() - st)
    dx1,dy1 = decode_one_voxels(dense_pts[:,5])
    dy2,dx2 = decode_one_voxels(dense_pts[:,6])
    xs = dense_pts[:,2] * vs[0] + pcr[0]
    ys = dense_pts[:,3] * vs[1] + pcr[1]
    pts = np.stack([dx1,dy1,dx2,dy2],-1).reshape(-1, 2, 2)
    pts *= vs[0]
    pts[...,0] += xs.reshape(-1,1)
    pts[...,1] += ys.reshape(-1,1)
    # axes = init_axes(dpi=2000)
    # draw_multi_grid_range(axes)
    # draw_multi_grid_range(axes, grid_size=0.8, bv_range=(-30.4, -15.2, 30.4, 15.2), rgb='g')
    # for k, l in enumerate(pts):
    #     # if dense_pts[k,0] != 0:
    #     #     break
    #     sc = int(dense_pts[k][-1])
    #     if k > 10000:
    #         break
    #     # print(l.shape)
    #     # l = np.stack(l,0)
    #     # if len(l) < 3:
    #         # continue
    #     # print(k)
    #     axes = matplot_lines(l[:,0], l[:,1], rgb=rgb[sc], scale=0.05, axes=axes)
    # plt.show()
if __name__ == "__main__":
    test()
